'''
使用方法
1、下载commonvoice 解压
2、复制VoiceToImg.py到解压过的目录
3、运行CommonVoiceSplit() 将commonvoice中的mp3文件按人分组
4、解压后的数据量比较大，这里学习使用，仅保留20人的数据其余的删除掉，我这里找的都是有5段mp3语音的文件夹（注，文件夹要重命名01~20）
5、运行ChangeMp3ToWavToImg() 将 分组后的mp3批量生成固定尺寸的语谱图
'''
import os
import shutil
from pydub import AudioSegment
import librosa  
import numpy as np
import matplotlib.pyplot as plt
import cv2

'''
把common voice数据集中的语音按人分组
将VoiceToImg.py放在commonvoice下载后解压的目录下
运行 CommonVoiceSplit()
将validated.tsv中标注的语音按照人的顺序号分组待后续处理
'''
def CommonVoiceSplit():
    rootpathout = '按人分组\\'
    filename = 'validated.tsv'
    f = open(filename, 'rb')
    content = f.read()
    str_content = content.decode('utf-8')
    linespit = str_content.split('\n')
    ids = {}
    for i in range(1 , len(linespit)):
        line = linespit[i].replace('  ', ' ')
        while '\t' in line:
            line = line.replace('\t', ' ')
        while '  ' in line:
            line = line.replace('  ', ' ')
        items = line.split(' ')
        if len(items) < 2:
            continue
        orgFile = 'clips\\' + items[1]
        if not os.path.exists(orgFile):
            continue
        if items[0] not in ids:
            ids[items[0]] = 0
        ids[items[0]] = ids[items[0]] + 1
        index = list(ids.keys()).index(items[0])
        groupPath = rootpathout + '\\' + str(index).zfill(4)
        if not os.path.isdir(groupPath):
            os.makedirs(groupPath)
        targetFile = groupPath + '\\' + items[1]
        if not os.path.exists(targetFile):
            shutil.copyfile(orgFile, targetFile)
            
'''
把声波序列生成对应的语谱图并保存
groupNp 声波序列
NFFT
framerate 采样率
framesize 抽样数
overlapSize 帧移数量
exportFile 保存文件名
index 图片索引 最终保存 exportFile + index + .png
'''
def ExportImg(groupNp, NFFT, framerate, framesize, overlapSize, exportFile, index):
    spectrum,freqs,ts,fig = plt.specgram(groupNp,
                                         NFFT = NFFT,
                                         Fs = framerate,
                                         window=np.hanning(M = framesize),
                                         noverlap=overlapSize,
                                         mode='default',
                                         scale_by_freq=True,
                                         sides='onesided',
                                         scale='dB',
                                         xextent=None)#绘制频谱图
    
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)
    plt.margins(0,0)
    plt.axis('off')
    plt.savefig("tmp.png")  
    '''   
    plt.ylabel('Frequency')
    plt.xlabel('Time')
    plt.title("Spectrogram")
    plt.savefig("tmp1.png")  
    '''   
    plt.show()
    src = cv2.imread("tmp.png")
    reshape = np.shape(spectrum)
    imageChange = cv2.resize(src, (reshape[1], reshape[0]))
    #imageChange = cv2.cvtColor(imageChange, cv2.COLOR_RGB2GRAY)
    #rgb  rg=黄色  判断r b的大小来区分前景背景
    cv2.imencode('.png', imageChange)[1].tofile(exportFile + str(index) + '.png')
     
'''
wav文件生成语谱图
对wav数据筛选，对声音强度高的1.5S截取生成声纹图  对强度低的 持续1.5S截取生成背景噪音语谱图
nplist wav序列
exportFile 声纹图目录
index 声纹图索引
backFile 背景语谱图目录
backIndex 背景谱图索引
framelength 帧时长 默认 0.025秒
framerate wav采样率 默认11025
'''       
def WavToImg(nplist, exportFile, index, backFile, backIndex, framelength = 0.025, framerate = 11025):
    framesize = framelength*framerate #每帧点数 N = t*fs,通常情况下值为256或512,要与NFFT相等\
                                    #而NFFT最好取2的整数次方,即framesize最好取的整数次方
 
    #找到与当前framesize最接近的2的正整数次方
    nfftdict = {}
    lists = [32,64,128,256,512,1024]
    for i in lists:
        nfftdict[i] = abs(framesize - i)
    sortlist = sorted(nfftdict.items(), key=lambda x: x[1])#按与当前framesize差值升序排列
    framesize = int(sortlist[0][0])#取最接近当前framesize的那个2的正整数次方值为新的framesize
     
    NFFT = framesize #NFFT必须与时域的点数framsize相等，即不补零的FFT
    overlapSize = 1.0/2 * framesize #重叠部分采样点数overlapSize约为每帧点数的1/3~1/2
    overlapSize = int(round(overlapSize))#取整
    print("帧长为{},帧叠为{},傅里叶变换点数为{}".format(framesize,overlapSize,NFFT))
    groupFrameCount = 6 #1S
    for i in range(framesize // 10):
        listsub = nplist[i * 10 : ]
        
        #每帧的最大值
        listMax = []        
        for i in range(len(listsub) // framesize):
            listMax.append(np.max(listsub[i * framesize : (i + 1) * framesize]))
        #整段最大值
        totalMax = np.max(listMax)
        #归一化
        listMax = listMax / totalMax
        #10帧一组 最大值超过0.1的个数比率
        listRate = []
        for l in range(len(listMax) // 10 - 1):
            count = 0
            for j in range(10): #2S
                if listMax[l * 10 + j] > 0.1:
                    count += 1
            listRate.append(count / 10.0)
        inGroup = 0
        for j in range(len(listRate) - groupFrameCount):
            if inGroup > 0:
                inGroup -= 1
                continue
            isTrue = True
            #连续8帧大于0.2
            for k in range(groupFrameCount):
                if listRate[j + k] < 0.3:
                    isTrue = False
                    break
            if isTrue:
                inGroup = groupFrameCount * 0.8
                groupNp = listsub[framesize * 10 * j : framesize * 10 * (j + groupFrameCount)]
                ExportImg(groupNp, NFFT, framerate, framesize, overlapSize, exportFile, index)
                index += 1
         
        inGroup = 0
        for j in range(len(listRate) - groupFrameCount):
            if inGroup > 0:
                inGroup -= 1
                continue
            isTrue = True
            #连续8帧大于0.2
            for k in range(groupFrameCount):
                if listRate[j + k] > 0.1:
                    isTrue = False
                    break
            if isTrue:
                inGroup = groupFrameCount
                groupNp = listsub[framesize * 10 * j : framesize * 10 * (j + groupFrameCount)]
                ExportImg(groupNp, NFFT, framerate, framesize, overlapSize, backFile, backIndex)
                backIndex += 1       
    return index, backIndex

'''
按人分组以后的commonvoice MP3批量生成对应的语谱图，用于后续制作数据集
'''
def ChangeMp3ToWavToImg():
    pathName = '按人分组\\'
    backPath = '按人分组\\0000\\'
    if not os.path.isdir(backPath):
        os.makedirs(backPath)
    #清理
    for file in os.listdir(pathName):
        img_path = pathName + file  #每类图片的地址
        for imgfile in os.listdir(img_path):
            if imgfile.endswith('png') or imgfile.endswith('wav') :
                img_file_path = img_path + '\\' + imgfile
                os.remove(img_file_path)
    backIndex = 0
    for file in os.listdir(pathName):
        img_path = pathName + file  #每类图片的地址
        imageIndex = 0
        for imgfile in os.listdir(img_path):
            if imgfile.endswith('mp3'):
                img_file_path = img_path + '\\' + imgfile
                print('changing file ' + img_file_path)
                sound = AudioSegment.from_mp3(img_file_path)
                tar_file = img_file_path + '.wav'
                sound.export(tar_file, format="wav")
                y, s = librosa.load(tar_file, sr=11025) # Downsample 44.1kHz to 8kHz
                imageIndex, backIndex = WavToImg(y, tar_file, imageIndex, backPath, backIndex)
    
    